import subprocess
from collections import OrderedDict

import argparse
import os
import warnings
import numpy as np

import mmcv
import torch
from mmcv import Config

from mmdet.datasets import build_dataloader, build_dataset, replace_ImageToTensor


def parse_args():
    parser = argparse.ArgumentParser(description="MMDet test (and eval) a model")
    parser.add_argument("config", help="test config file path")
    parser.add_argument("--out_dir", help="output directory")
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    cfg.data.test.ann_file = "/root/work/iva_detection/jit_model/calib_data.json"
    cfg.data.test.img_prefix = "/root/work/data/phcl"
    cfg.data.test.test_mode = True
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(
        dataset, samples_per_gpu=1, workers_per_gpu=1, dist=False, shuffle=False
    )

    for i, info in enumerate(data_loader):
        image = info["img"][0].squeeze(0).numpy()
        print(image.shape)
        image.astype(np.float32).flatten().tofile(
            os.path.join(args.out_dir, f"{i}.bin")
        )
    file_list = [f"{i}.bin" for i in range(len(data_loader))]
    with open(os.path.join(args.out_dir, "file_list"), "w") as f:
        for image_path in file_list:
            f.write("{}\n".format(image_path))


if __name__ == "__main__":
    main()
